"""
    Export line detections and descriptors given a list of input images.
"""
import os
import argparse
import cv2
import numpy as np
import torch
from tqdm import tqdm
import gc

from experiment import load_config
from model.line_detector import LineDetector


def export_descriptors(images_list, ckpt_path, config, device, extension,
                       output_folder):
    # Extract the image paths
    with open(images_list, 'r') as f:           #以只读方式打开文件
        image_files = f.readlines()             #读取所有行并返回列表
    #path.strip('\n')同时删除字符串左右两边的空格
    image_files = [path.strip('\n') for path in image_files]

    line_detection = LineDetector(
        config["model_cfg"], ckpt_path, device, config["line_detector_cfg"])
    print("\t Successfully initialized model")

    # Run the inference on each image and write the output on disk
    #tqdm()显示进度条
    for img_path in tqdm(image_files):
        img = cv2.imread(img_path, 0)           #灰度图的方式载入
        # if img.shape[1]>1000:
        #     img = cv2.resize(img,dsize=(1280,720))
        img = cv2.resize(img, dsize=(320, 320))
        img = torch.tensor(img[None, None] / 255., dtype=torch.float,
                           device=device)

        # Run the line detection and description
        # {"descriptor","heatmap","junctions","line_segments"}
        # batch x H/4 x W/4 x 128, H x W, num_junc x 2, num_segs x 2 x 2
        ref_detection = line_detection(img,return_heatmap=True)
        ref_line_seg = ref_detection["line_segments"]   #num_segs x 2 x 2
        ref_heat_map = ref_detection["heatmap"]
        ref_heat_map_raw = ref_detection["heatmap_raw"]
        ref_angle = ref_detection["angle"]
        # import matplotlib.pyplot as plt
        # plt.figure("raw")
        # plt.imshow(ref_heat_map_raw)
        # plt.figure("heatmap")
        # plt.imshow(ref_heat_map)
        # plt.show()
        # import ipdb; ipdb.set_trace()

        # Write the output on disk
        os.makedirs(output_folder, exist_ok=True)
        #os.path.basename(path)返回path最后的文件名,os.path.splitext()分离文件名与扩展名
        img_name = os.path.splitext(os.path.basename(img_path))[0]           #取不带扩展名的文件名
        output_file = os.path.join(output_folder, img_name + extension)      #拼接文件路径
        #以压缩的.npz格式将多个数组保存到一个文件中
        #np.savez_compressed(file,arr)    file:保存数据的文件名，.npz将被附加到文件名后面
        #arr：要保存到文件的数组，使用关键字参数为数组分配名称
        np.savez_compressed(output_file, line_seg=ref_line_seg,
                            heatmap=ref_heat_map, heatmap_raw = ref_heat_map_raw, angle = ref_angle)
        gc.collect()
        torch.cuda.empty_cache()

if __name__ == "__main__":
    # Parse input arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--img_list", type=str, required=True,
                        help="List of input images in a text file.")
    parser.add_argument("--output_folder", type=str, required=True,
                        help="Path to the output folder.")
    parser.add_argument("--config", type=str,
                        default="config/export_line_features.yaml")
    parser.add_argument("--checkpoint_path", type=str,
                        default="pretrained_models/sold2_wireframe.tar")
    parser.add_argument("--extension", type=str, default=None)
    args = parser.parse_args()

    # Get the device
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # Get the model config, extension and checkpoint path
    config = load_config(args.config)       #python字典
    #获取文件绝对路径
    ckpt_path = os.path.abspath(args.checkpoint_path)
    #如果args.extension是None则extension是'sold2'，否则是args.extension
    extension = 'sold2' if args.extension is None else args.extension
    extension = "." + extension

    export_descriptors(args.img_list, ckpt_path, config, device, extension,
                       args.output_folder)
